'''
Created on Aug 19, 2009

@author: mkiyer
'''

from bx.intervals.intersection import Interval, IntervalTree

class GenomicInterval(object):    
    def __init__(self, chrom, start, end):
        self.chrom = chrom
        self.start = start
        self.end = end

    def parse_line(self, line):
        fields = line.strip().split('\t')
        self.chrom = fields[0]
        self.start = int(fields[1])
        self.end = int(fields[2])

class BedInterval(GenomicInterval):
    def __init__(self, chrom, start, end, name='.', score=1000, strand='+', 
                 tx_start=0, tx_end=0, item_rgb=0, 
                 block_count=0, block_sizes=None, block_starts=None):
        super(BedInterval, self).__init__(start, end, chrom=chrom, strand=strand)
        self.name = name
        self.score = score
        self.tx_start = tx_start
        self.tx_end = tx_end
        self.item_rgb = item_rgb
        self.block_count = block_count
        self.block_sizes = block_sizes
        self.block_starts = block_starts
    
    def parse_line(self, line):
        fields = line.strip().split('\t')
        self.chrom = fields[0]
        self.start = int(fields[1])
        self.end = int(fields[2])
        self.name = fields[3]
        self.score = int(fields[4])
        self.strand = fields[5]
        self.tx_start = int(fields[6])
        self.tx_end = int(fields[7])
        self.item_rgb = fields[8]
        self.block_count = int(fields[9])
        self.block_sizes = map(int, fields[10].split(','))
        self.block_starts = map(int, fields[11].split(','))

def parse_raw_interval_file(fhd, category):
    for id, line in enumerate(fhd):
        fields = line.strip().split('\t')
        chrom = fields[0]
        start = int(fields[1])
        end = int(fields[2])
        yield Interval(start, end, value=('%s_%d' % (category, id), category), chrom=chrom)
        
def parse_bed_file(fhd, category):
    for line in fhd:
        fields = line.strip().split('\t')
        chrom = fields[0]
        start = int(fields[1])
        end = int(fields[2])
        name = fields[3]
        strand = fields[5]
        yield Interval(start, end, value=(name, category), chrom=chrom, strand=strand)

def parse_kgtxinfo(fhd):
    kgtxinfo = {}
    for line in fhd:
        fields = line.strip().split('\t')
        kgtxinfo[fields[0]] = fields[1]
    return kgtxinfo

def parse_gene_pred(fhd, category):
    '''
    name varchar(255) NOT NULL default '',
    chrom varchar(255) NOT NULL default '',
    strand char(1) NOT NULL default '',
    txStart int(10) unsigned NOT NULL default '0',
    txEnd int(10) unsigned NOT NULL default '0',
    cdsStart int(10) unsigned NOT NULL default '0',
    cdsEnd int(10) unsigned NOT NULL default '0',
    exonCount int(10) unsigned NOT NULL default '0',
    exonStarts longblob NOT NULL,
    exonEnds longblob NOT NULL,
    proteinID varchar(40) NOT NULL default '',
    alignID varchar(255) NOT NULL default '',
    '''
    for line in fhd:
        fields = line.strip().split('\t')
        name = fields[0]
        chrom = fields[1]
        strand = fields[2]
        txstart = int(fields[3])
        txend = int(fields[4])
        yield Interval(txstart, txend, value=(name, category), chrom=chrom, strand=strand)

def load_ucsc_genes(kgfhd, kgtxinfofhd):
    kgtxinfodict = parse_kgtxinfo(kgtxinfofhd)
    # knowngene file is genePred format
    for interval in parse_gene_pred(kgfhd, 'ucsc'):
        ucsc_category = kgtxinfodict[interval.value[0]]
        if ucsc_category == "coding":
            interval.value = (interval.value[0], 'ucsc_coding')
        else:
            interval.value = (interval.value[0], 'ucsc_noncoding')
        #interval.value = (interval.value[0], 'ucsc_%s' % kgtxinfodict[interval.value[0]])
        yield interval
    del kgtxinfodict
        
